from crewai.tools import BaseTool
from typing import Type, List, Optional
from pydantic import BaseModel, Field
import requests
import json
import csv
import os


class ETFRequest(BaseModel):
    """单个ETF基金请求参数"""
    code: str = Field(..., description="基金代码（必填字段），格式如：sh516970（上海）、sh512480（上海）、sz159570（深圳）等。注意：上海基金代码以'sh'开头，深圳基金代码以'sz'开头")
    limit: Optional[int] = Field(None, description="返回历史K线数据的条数（可选参数）。如果指定limit，则返回最近N条历史数据；如果不指定，则自动限制为750条。最大支持750条。例如：limit=500表示返回最近500条K线数据")


class ETFDataToolInput(BaseModel):
    """ETF数据工具输入参数"""
    etf_requests: List[ETFRequest] = Field(..., description="ETF基金请求列表（必填字段）。这是一个列表，包含一个或多个基金查询请求。每个请求是一个字典对象，必须包含'code'字段（基金代码），可选包含'limit'字段（数据条数限制）。注意：整个参数应该是一个包含单个'etf_requests'字段的对象，而不是数组。例如：{'etf_requests': [{'code': 'sz159570', 'limit': 500}]}")


class ETFDataTool(BaseTool):
    name: str = "etf_data_tool"
    description: str = (
        "【ETF基金数据查询工具】\n"
        "功能：获取ETF基金的最新净值、累计净值、日增长率以及历史K线数据。支持同时查询多个基金，支持指定返回的历史K线数据条数。\n\n"
        "核心特性：\n"
        "1. 支持单个或多个基金查询：可一次查询多个不同的ETF基金\n"
        "2. 灵活的数据量控制：可以指定返回的历史K线数据条数，最大支持750条。未指定limit或超过750时自动限制为750条\n"
        "3. 完整的数据返回：包含最新净值、累计净值、日增长率以及完整的K线数据（开盘、收盘、最高、最低、成交量等）\n"
        "4. **自动保存CSV文件**：获取数据后，会自动将历史K线数据保存为CSV文件到 `out/` 目录，文件名格式为 `{基金代码}_historical.csv`（例如：sh512690_historical.csv）。CSV文件包含完整的K线数据字段，便于后续量化分析使用\n\n"
        "基金代码格式：\n"
        "- 上海证券交易所ETF：以'sh'开头，如 sh516970、sh512480\n"
        "- 深圳证券交易所ETF：以'sz'开头，如 sz159570\n\n"
        "参数结构说明：\n"
        "工具接受一个包含'etf_requests'字段的对象，'etf_requests'是一个列表，每个元素是一个基金查询请求。\n\n"
        "正确的参数格式（JSON对象，不是数组）：\n"
        "{\n"
        "  \"etf_requests\": [\n"
        "    {\"code\": \"sz159570\", \"limit\": 500},\n"
        "    {\"code\": \"sh516970\", \"limit\": 250}\n"
        "  ]\n"
        "}\n\n"
        "错误格式示例（不要使用）：\n"
        "- 数组格式：[{\"etf_requests\": [...]}]  ❌ 错误\n"
        "- 多个etf_requests对象：[{\"etf_requests\": [...]}, {\"etf_requests\": [...]}]  ❌ 错误\n\n"
        "使用示例：\n"
        "1. 查询单个基金，获取500条历史数据：\n"
        "   {\"etf_requests\": [{\"code\": \"sz159570\", \"limit\": 500}]}\n"
        "2. 查询多个基金，第一个指定条数，第二个获取全部：\n"
        "   {\"etf_requests\": [{\"code\": \"sz159570\", \"limit\": 500}, {\"code\": \"sh516970\"}]}\n"
        "3. 查询单个基金，未指定limit时自动限制为750条：\n"
        "   {\"etf_requests\": [{\"code\": \"sz159570\"}]}\n\n"
        "返回内容：对于每个基金，返回基金代码、名称、最新净值、累计净值、日增长率，以及历史K线数据（包含日期、开盘价、收盘价、最高价、最低价、成交量、涨跌额、涨跌幅、振幅等信息）。\n\n"
        "**重要提示**：历史K线数据会自动保存为CSV文件到 `out/` 目录，文件名格式为 `{基金代码}_historical.csv`。\n"
        "例如：查询基金代码为 sh512690 的数据后，CSV文件会保存为 `out/sh512690_historical.csv`。\n"
        "量化分析师可以使用文件读取工具直接读取该CSV文件进行量化分析，避免因上下文数据过大导致的问题。"
    )
    args_schema: Type[BaseModel] = ETFDataToolInput

    def _save_to_csv(self, fund_data: dict, output_dir: str = "out") -> str:
        """
        将基金的K线数据保存为CSV文件
        
        Args:
            fund_data: 单个基金的数据字典
            output_dir: 输出目录，默认为"out"
            
        Returns:
            str: CSV文件的路径，如果保存失败则返回空字符串
        """
        try:
            # 确保输出目录存在
            os.makedirs(output_dir, exist_ok=True)
            
            # 获取基金代码
            fund_code = fund_data.get("code", "unknown")
            
            # 构建CSV文件路径
            csv_filename = f"{fund_code}_historical.csv"
            csv_path = os.path.join(output_dir, csv_filename)
            
            # 获取K线数据
            if "kData" not in fund_data or not isinstance(fund_data["kData"], list):
                return ""
            
            klines = fund_data["kData"]
            if not klines:
                return ""
            
            # 写入CSV文件（使用 LF 行尾符以确保跨平台一致性）
            with open(csv_path, 'w', newline='', encoding='utf-8') as csvfile:
                fieldnames = ['日期', '开盘价', '最高价', '最低价', '收盘价', '成交量', '涨跌额', '涨跌幅', '振幅']
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                
                writer.writeheader()
                for kline in klines:
                    if isinstance(kline, dict):
                        writer.writerow({
                            '日期': kline.get('d', ''),
                            '开盘价': kline.get('o', ''),
                            '最高价': kline.get('h', ''),
                            '最低价': kline.get('l', ''),
                            '收盘价': kline.get('c', ''),
                            '成交量': kline.get('v', ''),
                            '涨跌额': kline.get('zde', ''),
                            '涨跌幅': kline.get('zd', ''),
                            '振幅': kline.get('zf', '')
                        })
            
            # 确保文件使用 LF 行尾符（统一为 Unix 格式）
            with open(csv_path, 'rb') as f:
                content = f.read()
            # 替换所有 CRLF 为 LF，并移除单独的 CR
            content = content.replace(b'\r\n', b'\n').replace(b'\r', b'\n')
            with open(csv_path, 'wb') as f:
                f.write(content)
            
            return csv_path
            
        except Exception as e:
            # 保存失败时返回空字符串，不影响主流程
            print(f"Warning: Failed to save CSV file: {e}")
            return ""

    def _format_result_as_text(self, result: dict) -> str:
        """
        将API返回的JSON数据格式化为易于大模型理解的文本格式
        
        Args:
            result: API返回的JSON数据
            
        Returns:
            str: 格式化的文本数据
        """
        try:
            formatted_text = "ETF基金数据查询结果:\n\n"
            
            # 检查是否有错误
            if "error" in result:
                return f"查询失败: {result['error']}"
            
            # 检查是否有数据
            if "data" not in result:
                return "未获取到有效数据"
            
            data = result["data"]
            
            # 如果是列表，遍历每个基金的数据
          
            for i, fund_data in enumerate(data, 1):
                formatted_text += f"基金 {i}:\n"
                
                # 基金基本信息
                if "code" in fund_data:
                    formatted_text += f"  基金代码: {fund_data['code']}\n"
                if "name" in fund_data:
                    formatted_text += f"  基金名称: {fund_data['name']}\n"
                if "valueDate" in fund_data:
                    formatted_text += f"  净值日期: {fund_data['valueDate']}\n"
                if "value" in fund_data:
                    formatted_text += f"  净值: {fund_data['value']}\n"
                if "totalValue" in fund_data:
                    formatted_text += f"  累计净值: {fund_data['totalValue']}\n"
                if "dailyRate" in fund_data:
                    formatted_text += f"  日增长率: {fund_data['dailyRate']}%\n"
                
                # 保存K线数据为CSV文件
                csv_path = self._save_to_csv(fund_data)
                if csv_path:
                    formatted_text += f"  CSV文件已保存: {csv_path}\n"
                
                # K线数据
                if "kData" in fund_data and isinstance(fund_data["kData"], list):
                    klines = fund_data["kData"]
                    formatted_text += f"  历史K线数据 (共{len(klines)}条):\n"
                    
                    for j, kline in enumerate(klines, 1):  
                        if isinstance(kline, dict):
                            formatted_text += f"    第{j}条: "
                            if "d" in kline:
                                formatted_text += f"日期:{kline['d']}, "
                            if "o" in kline:
                                formatted_text += f"开盘:{kline['o']}, "
                            if "h" in kline:
                                formatted_text += f"最高:{kline['h']}, "
                            if "l" in kline:
                                formatted_text += f"最低:{kline['l']}, "
                            if "c" in kline:
                                formatted_text += f"收盘:{kline['c']}, "
                            if "v" in kline:
                                formatted_text += f"成交量:{kline['v']}"
                            if "zde" in kline:
                                formatted_text += f"涨跌额:{kline['zde']}, "
                            if "zd" in kline:
                                formatted_text += f"涨跌幅:{kline['zd']}%, "
                            if "zf" in kline:
                                formatted_text += f"振幅:{kline['zf']}%, " 
                            formatted_text += "\n"
                
                formatted_text += "\n"
            
            return formatted_text
            
        except Exception as e:
            # 如果格式化失败，返回原始JSON
            return f"数据格式化失败: {str(e)}\n原始数据: {json.dumps(result, ensure_ascii=False, indent=2)}"

    def _run(self, etf_requests: List[dict]) -> str:
        """
        调用API接口获取ETF基金数据
        
        Args:
            etf_requests: ETF基金请求列表，每个元素是包含'code'和'limit'字段的字典
            
        Returns:
            str: 格式化的文本数据
        """
        try:
            # API接口地址
            api_url = "https://tools.zhangjianyong.top/api/etf/getData"
            
            # 构建请求数据
            request_data = []
            MAX_LIMIT = 750  # 最大历史K线数据条数限制
            for etf_req in etf_requests:
                # etf_req 实际是字典类型，使用字典访问方式
                req_item = {"code": etf_req["code"]}
                limit = etf_req.get("limit")
                # 限制最大条数为750，未指定或超过750时自动设为750
                if limit is None or not isinstance(limit, int) or limit > MAX_LIMIT or limit <= 0:
                    req_item["limit"] = MAX_LIMIT
                else:
                    req_item["limit"] = limit
                request_data.append(req_item)
            
            # 发送POST请求
            headers = {
                'Content-Type': 'application/json',
                'User-Agent': 'ETF-Data-Tool/1.0'
            }
            
            response = requests.post(
                api_url, 
                json=request_data, 
                headers=headers,
                timeout=30
            )
            
            # 检查响应状态
            response.raise_for_status()
            
            # 返回格式化的文本数据，便于大模型理解
            result = response.json()
            return self._format_result_as_text(result)
            
        except requests.exceptions.RequestException as e:
            error_msg = f"API请求失败: {str(e)}"
            return f"ETF数据查询失败: {error_msg}"
            
        except json.JSONDecodeError as e:
            error_msg = f"JSON解析失败: {str(e)}"
            return f"ETF数据解析失败: {error_msg}"
            
        except Exception as e:
            error_msg = f"未知错误: {str(e)}"
            return f"ETF数据查询出现未知错误: {error_msg}"


def main():
    """测试ETF数据工具的主函数"""
    print("=== ETF数据工具测试 ===")
    
    # 创建工具实例
    etf_tool = ETFDataTool()
    
    # 测试用例1：查询单个基金（不指定条数）
    print("\n1. 测试单个基金查询（不指定条数）:")
    test_requests_1 = [
        ETFRequest(code="sh516970")
    ]
    
    try:
        result_1 = etf_tool._run(test_requests_1)
        print("请求数据:", [{"code": req.code, "limit": req.limit} for req in test_requests_1])
        print("响应结果:")
        print(result_1)
    except Exception as e:
        print(f"测试1失败: {e}")
    
    # 测试用例2：查询多个基金（指定条数）
    print("\n2. 测试多个基金查询（指定条数）:")
    test_requests_2 = [
        ETFRequest(code="sh516970"),
        ETFRequest(code="sh512480", limit=10)
    ]
    
    try:
        result_2 = etf_tool._run(test_requests_2)
        print("请求数据:", [{"code": req.code, "limit": req.limit} for req in test_requests_2])
        print("响应结果:")
        print(result_2)
    except Exception as e:
        print(f"测试2失败: {e}")
    
   
    
    print("\n=== 测试完成 ===")


if __name__ == "__main__":
    main()
